#pragma once
 
#include <maya/MArgList.h>
#include <maya/MObject.h>
#include <maya/MGlobal.h>
#include <maya/MPxCommand.h>
#include <maya/MFloatPointArray.h>
#include <maya/MItDag.h>
#include <maya/MDagPath.h>
#include <maya/MFnMesh.h>
#include <maya/MFloatPoint.h>
#include <maya/MFloatVector.h>
#include <maya/MColor.h>
#include <maya/MFn.h>
#include <maya/MThreadPool.h>
#include "BBox.h"
#include "VoxelSpace.h"
#include "MeshInfo.h"
#include <limits>
#include <vector>
#include <fstream>
#include <iostream>
#include <map>

class Material;
extern const BBox BBoxSentinel;
class RayTracer; 

#define CHI_INV_COUNT 100

struct RenderPlaneInfo
{
  RenderPlaneInfo()
    : ImgResX(0), ImgResY(0), LeftBottom(MPoint()), PixelHeight(0.0), PixelWidth(0.0),
    eyePoint(MPoint()), viewDirection(MVector(0,0,0)), upDirection(MVector(0,0,0)), rightDirection(MVector(0,0,0)),
    isOrtho(false), ImgPlaneH(0.0), ImgPlaneW(0.0), aspectRatio(1.0)
  {}

  int ImgResX;
  int ImgResY;
  MPoint LeftBottom;
  double PixelHeight;
  double PixelWidth;
  MFloatPoint eyePoint;
  MVector viewDirection;
  MVector upDirection;
  MVector rightDirection;
  bool isOrtho;
  double ImgPlaneH;
  double ImgPlaneW;
  double aspectRatio;
};

struct Range
{
  Range(int s=-1, int e=-1) : y_start(s), y_end(e) {}
  int y_start;
  int y_end;
};

struct PixelSample
{
  PixelSample()
    : imgX(-1)
    , imgY(-1)
    , outColor(MColor(0.0f,0.0f,0.0f))
    , raySource(MFloatPoint(0.0f,0.0f,0.0f))
    , rayDir(0.0f, 0.0f, 0.0f) 
  {}

  int           imgX;
  int           imgY;
  MColor        outColor;
  MFloatPoint   raySource;
  MFloatVector  rayDir;
  // ... other sample params
};


struct LightInfo
{
  MDagPath      light_path;
  MFn::Type     light_type;
  MFloatPoint   pos;
  MFloatVector  dir;
  MColor        color;
  float         intensity;
};

enum SuperSampleType
{
  SuperSampleType_Unknown=-1,
  SuperSampleType_None=0,
  SuperSampleType_Uniform,
  SuperSampleType_Jittered,
  SuperSampleType_Random,
  SuperSampleType_Adaptive,
  SuperSampleType_numerof
};

const MString SST_Uniform ("uniform");
const MString SST_Jittered ("jittered");
const MString SST_Random ("random");
const MString SST_Adaptive ("adaptive");
MString SST2STR(SuperSampleType sst);

struct ThreadStatData
{
  ThreadStatData() 
    : total_rays(0) 
    , total_samples(0)
  {}

  std::vector<float> pixel_times;
  int total_rays;
  int total_samples;
  std::vector<int> pixel_samples;
  CIStats ci_stats;
  std::vector<int> sample_ray_depths;
};

// the MThreadCallback func
struct ThreadData
{
  ThreadData(){}
  int             tidx;
  Range           chunk_range;
  RayTracer*      ray_tracer;
  ThreadStatData* pstat_data;
  // timers
};

struct TaskData
{
  TaskData(RayTracer* rt, const std::vector<Range>& ranges_) 
  {
    ranges = ranges_;
    thread_count = (int)ranges.size();
    for(int i=0; i<thread_count; ++i){
      stats.push_back(ThreadStatData());
    }
    ray_tracer = rt;
  }

  int                         thread_count;
  std::vector<Range>          ranges;
  std::vector<ThreadStatData> stats;
  RayTracer*                  ray_tracer;
};

void DecomposeRayTrace(void *data, MThreadRootTask *root);
MThreadRetVal RayTraceChunk(void *data);

class RayTracer : public MPxCommand 
{

public:

  RayTracer();
  virtual MStatus   doIt(const MArgList& argList);
  static bool       intArg(const MArgList& args, unsigned int &indx, int & res);
  static bool       intArrayArg(const MArgList& args, unsigned int &indx, unsigned int count, std::vector<int> & res);
  static bool       fltArg(const MArgList& args, unsigned int &indx, float & res);
  static bool       strArg(const MArgList& args, unsigned int &indx, MString &res);

  static MStatus    parseArgs(const MArgList& args, int& width, int& height, int& voxelFactor, int& superSample, SuperSampleType& sst,
    int& samplingRate, float& tolerance, int& misr, int& masr, int& rayDepth, int& verbose, int& draw_x, int& draw_y, bool& drawOnly, int& threadCount);
  
  static void*      creator();
  static bool       getBoundingBox(const MDagPath& dagPath, bool quiet, BBox& bbox_);
  static bool       getBoundingBox(const std::vector<BBox>& bboxes, BBox& bbox_);
  
  static MColor     calcColorAtPt(VoxelSpace& vxspc, float uv1, float uv2, 
    const MFloatPoint& vertPt, const MFloatVector& vertNorm, const MFloatPoint& eyePoint, 
    const Material& material, const std::vector<LightInfo>& light_info_list, 
    int minMeshIdx, int minHitFace, float minHitBary1, float minHitBary2,
    CIStats& ci_stat,
    const MColor& ref_color, const MColor& trans_color, float fres_kr);

  static bool				lightSeesPoint(const VoxelSpace& vxspc, const MFloatPoint& light_pos, const MVector& light_vec, 
    int mesh_idx, int face_idx, float bary1, float bary2, 
    CIStats& ci_stat,
    bool directional_light = false);
  static MString	  buildOutputDirectoryName();
  static void       triangulateMesh(const MFnMesh& mesh);
  static void       triangulate_meshes();
  void              create_meshinfo_list(std::vector<MeshInfo>& mesh_info_vec);
  void              create_triangles_list(std::vector<std::vector<triangle>>& triangles_vector, int& triangle_count, bool verbose);
  void              get_scene_bounding_box(BBox& scene_bb, float scale_factor=1.0);
  MStatus           get_renderplane_info(int width, int height, RenderPlaneInfo& rp_info);
  void              create_lightinfo_list(std::vector<LightInfo>& light_info_list);
  void				      computeChiInvForBetas();
  //void              cleanUpChiInv();

  bool              RayTrace(const MFloatPoint& raySource, const MFloatVector& rayDir, int traceDepth, float curr_refraction_idx, 
    MColor& out_color, 
    CIStats& ci_stats_, int& total_rays_, std::vector<int>& sample_ray_depths,
    bool draw, triangle_object* tri_skip=NULL);


  // params 
  int m_arg_width;
  int m_arg_height;
  int m_arg_voxelFactor;
  int m_arg_superSample;
  SuperSampleType m_arg_superSampleType;
  int m_arg_samplingRate; 
  float m_arg_tolerance;
  int m_arg_minSR;
  int m_arg_maxSR;
  int m_arg_rayDepth;
  int m_arg_verbose;
  int m_arg_draw_x;
  int m_arg_draw_y;
  bool m_arg_draw;
  bool m_arg_drawOnly;
  int m_arg_thread_count;

  double** m_chiInv_R;
  double beta_R,beta_G,beta_B;
  double averageSamplingRate;

  std::vector<MeshInfo> m_mesh_info_vec;
  std::vector<std::vector<triangle>> m_triangles_vector;
  RenderPlaneInfo m_rp_info;
  std::vector<LightInfo> m_light_info_list;
  MImage m_img;
  VoxelSpace m_vxspc;
  static Visualizer s_viz;
  int m_core_count;
};